import os
import torch
import logging
from typing import Dict, List, Any, Tuple, Optional
from pathlib import Path

# 从unified_representation导入训练函数
from ..unified_representation.train import train

logger = logging.getLogger(__name__)

class ModelTrainer:
    """
    ART模型训练类，负责模型配置与训练
    """
    
    def __init__(self, config: Dict[str, Any] = None):
        """
        初始化模型训练器
        
        参数:
            config: 模型配置参数
        """
        self.config = config or {}
        
        # 默认模型配置
        self.model_config = {
            "batch_size": self.config.get("batch_size", 32),       # 批处理大小，控制每次训练的样本数量
            "epochs": self.config.get("epochs", 1000),              # 训练轮数，模型训练的总迭代次数
            "num_heads": self.config.get("num_heads", 2),          # 注意力头数量，Transformer中的多头注意力参数
            "tf_layers": self.config.get("tf_layers", 1),          # Transformer层数，Transformer编码器的层数
            "gnn_hidden_dim": self.config.get("gnn_hidden_dim", 64), # 图神经网络隐藏层维度，控制GNN模型的复杂度
            "gnn_out_dim": self.config.get("gnn_out_dim", 32),     # 图神经网络输出层维度，控制GNN最终的输出特征维度
            "noise_rate": self.config.get("noise_rate", 0.3),      # 噪声率，用于dropout正则化，防止过拟合
            "gnn_layers": self.config.get("gnn_layers", 2),        # 图神经网络层数，GNN的层数越多，捕获的图结构特征越复杂
            "gru_hidden_dim": self.config.get("gru_hidden_dim", 32), # GRU隐藏层维度，控制序列特征提取的能力
            "gru_layers": self.config.get("gru_layers", 2),        # GRU层数，多层GRU可以提取更复杂的时序特征
            "learning_rate": self.config.get("learning_rate", 0.001), # 学习率，控制模型参数更新的步长大小
        }
    
    def adjust_num_heads(self):
        """
        确保num_heads能够整除instance_dim，避免PyTorch中出现"embed_dim must be divisible by num_heads"错误
        
        处理步骤:
        1. 检查instance_dim是否设置
        2. 检查instance_dim是否能被num_heads整除
        3. 如果不能整除，尝试调整num_heads为合适的值
        """
        # 步骤1: 检查instance_dim是否设置
        instance_dim = self.model_config.get('instance_dim', 0)
        if not instance_dim:
            return
            
        # 步骤2: 检查能否整除
        num_heads = self.model_config['num_heads']
        
        # 步骤3: 如果不能整除，调整num_heads
        if instance_dim % num_heads != 0:
            # 尝试找到合适的num_heads值
            for potential_heads in [2, 4, 8, 16, 1]:
                if instance_dim % potential_heads == 0 and potential_heads <= instance_dim:
                    old_heads = num_heads
                    self.model_config['num_heads'] = potential_heads
                    logger.info(f"将num_heads从 {old_heads} 调整为 {potential_heads} 以确保能够整除instance_dim ({instance_dim})")
                    break
    
    def update_model_config(self, channel_dim: int, instance_dim: int):
        """
        更新模型配置参数
        
        参数:
            channel_dim: 特征维度
            instance_dim: 实例维度
        """
        self.model_config['channel_dim'] = channel_dim
        self.model_config['instance_dim'] = instance_dim
        self.adjust_num_heads()
        logger.info(f"更新模型配置: channel_dim={channel_dim}, instance_dim={instance_dim}, num_heads={self.model_config['num_heads']}")
    
    def train_model(self, train_samples: List[Tuple]) -> Any:
        """
        训练ART模型
        
        参数:
            train_samples: 训练样本列表
            
        返回:
            训练好的模型
        """
        if not train_samples:
            logger.error("训练样本为空，无法训练模型")
            return None
        
        logger.info("开始训练ART模型")
        model = train(train_samples, self.model_config)
        logger.info("ART模型训练完成")
        
        return model
        
    def save_model(self, model: Any, output_path: Path, anomaly_id: str = None):
        """
        保存训练好的模型
        
        参数:
            model: 训练好的模型
            output_path: 输出路径
            anomaly_id: 异常ID，用于命名
            
        返回:
            str: 保存的模型路径
        """
        if model is None:
            logger.error("模型为空，无法保存")
            return None
            
        # 创建目录
        output_path.mkdir(parents=True, exist_ok=True)
        
        # 生成模型文件名
        model_name = f"art_model.pt"
        model_path = os.path.join(output_path, model_name)
        
        # 保存模型
        torch.save(model.state_dict(), model_path)
        logger.info(f"模型已保存到 {model_path}")
        
        return str(model_path) 